from utils.common import dump_json,load_logging_config
import os
import shutil
from load_acoustic_model import get_model_ins_from_json,get_network_json
import keras
import logging
from consta import path_sep

def add_acoustic_model():
    name='SpeechModel251BN1Small'
    typ='DL'
    filetype='source'
    content='structure'
    weights_file=''
    weights_type='' # ckpt,weights 见utils/common.py:load_weights_keras
    file='acoustic_models.keras_backend'
    classname='SpeechModel251BNSmall'
    if filetype=='source':
        if os.path.exists(file):
            if not os.path.abspath(file).startswith(os.path.abspath('acoustic_models')):
                target_acoustic_model_py=f'acoustic_models/{os.path.splitext(file.split("/")[-1])[0]}.py'
                if os.path.basename(target_acoustic_model_py) in os.listdir('acoustic_models'):
                    logging.error('声学模型模块重名')
                shutil.copy(file,target_acoustic_model_py)
                modulename=target_acoustic_model_py[:-3].replace('/','.')
            else:
                modulename=os.path.relpath(file,'.').replace(path_sep,'.')
                modulename=modulename if not modulename.endswith('.py') else modulename[:-3]
        elif file=='':
            modulename=f'acoustic_models.{name}'
        else:
            modulename=file if not file.endswith('.py') else file[:-3]
            modulename=modulename.replace('/','.')
    infer_classname='SpeechModel251BNSmall'
    infer_parameters='(1600,200,1),1429'
    infer_attribute='model_base'
    infer_input_layer_names=[]
    infer_output_layer_names=[] # https://blog.csdn.net/u011590738/article/details/126247488，构建推理模型的输入节点和输出节点，串起推理模型
    optimizer_file_path='tensorflow.keras.optimizers' # 按绝对路径处理，如果路径不存在，说明是已知的某个模块，如果是空字符串，说明是模型代码文件里的内容，如果真有对应的代码文件，将代码文件拷贝到optimizers目录下，并变换模块名为optimizers下的模块，见下面的操作
    if os.path.exists(optimizer_file_path):
        if not os.path.abspath(optimizer_file_path).startswith(os.path.abspath('optimizers')):
            target_optimizer_py=f'optimizers/{os.path.splitext(optimizer_file_path.split("/")[-1])[0]}.py'
            if os.path.basename(target_optimizer_py) in os.listdir('optimizers'): # 在复制文件前提前判断，提醒评测人员重名
                logging.error('优化器模块重名')
            shutil.copy(optimizer_file_path,target_optimizer_py)
            optimizer_module_name=target_optimizer_py[:-3].replace('/','.')
        else:
            optimizer_module_name=os.path.relpath(optimizer_file_path,'.').replace(path_sep,'.')
            optimizer_module_name=optimizer_module_name if not optimizer_module_name.endswith('.py') else optimizer_module_name[:-3]
    elif optimizer_file_path=='':
        optimizer_module_name=f'acoustic_models.{name}'
    else:
        optimizer_module_name=optimizer_file_path if not optimizer_file_path.endswith('.py') else optimizer_file_path[:-3]
        optimizer_module_name=optimizer_module_name.replace('/','.')
    optimizer_classname='Adam'
    optimizer_parameters='learning_rate=0.0001, beta_1=0.9, beta_2=0.999, weight_decay=0.0, epsilon=10e-8'
    optimizer_attribute=''
    optimizer={
        'modulename':optimizer_module_name,
        'classname':optimizer_classname,
        'parameters':optimizer_parameters,
        'attribute':optimizer_attribute
    }
    loss_file_path='acoustic_models/keras_backend'
    if os.path.exists(loss_file_path): # 如果loss_file_path用的是模型自身的损失函数，则loss_file_path=='!'，则loss_module_name=='!'
        if not os.path.abspath(loss_file_path).startswith(os.path.abspath('losses')):
            target_loss_py=f'losses/{os.path.splitext(loss_file_path.split("/")[-1])[0]}.py'
            if os.path.basename(target_loss_py) in os.listdir('losses'): # 在复制文件前提前判断，提醒评测人员重名
                logging.error('损失模块重名')
            shutil.copy(loss_file_path,target_loss_py)
            loss_module_name=target_loss_py[:-3].replace('/','.')
        else:
            loss_module_name=os.path.relpath(loss_file_path,'.').replace(path_sep,'.')
            loss_module_name=loss_module_name if not loss_module_name.endswith('.py') else loss_module_name[:-3]
    elif loss_file_path=='':
        loss_module_name=f'acoustic_models.{name}'
    else:
        loss_module_name=loss_file_path if not loss_file_path.endswith('.py') else loss_file_path[:-3]
        loss_module_name=loss_module_name.replace('/','.')
    loss_classname='sm251bn_loss'
    loss_parameters=''
    loss_attribute=''
    loss={
        'modulename':loss_module_name,
        'classname':loss_classname,
        'parameters':loss_parameters,
        'attribute':loss_attribute
    }
    parameters='(1600,200,1),1429'
    attribute='model'
    acoustic_model_json={
        'name':name,
        'type':typ,
        'filetype':filetype,
        'content':content,
        'weights_file':weights_file,
        'weights_type':weights_type,
        'file':file,
        'optimizer':optimizer,
        'loss':loss,
    }

    if filetype=='source':
        acoustic_model_json['modulename']=modulename
        acoustic_model_json['classname']=classname
        acoustic_model_json['parameters']=parameters
        acoustic_model_json['attribute']=attribute
    elif filetype=='framework_save':
        acoustic_model_json['file_ext']='' # 见utils/common.py:102 
        acoustic_model_json['framework']='tf.keras'
    
    acoustic_model_json['infer_classname']=infer_classname
    acoustic_model_json['infer_parameters']=infer_parameters
    acoustic_model_json['infer_attribute']=infer_attribute
    acoustic_model_json['infer_input_layer_names']=infer_input_layer_names
    acoustic_model_json['infer_output_layer_names']=infer_output_layer_names

    acoustic_model,_=get_model_ins_from_json(acoustic_model_json,'train')

    acoustic_model_json['network']=get_network_json(acoustic_model)

    if not f'{acoustic_model_json["name"]}.json' in os.listdir('jsons/acoustic_models'):
        dump_json(acoustic_model_json,f'jsons/acoustic_models/{acoustic_model_json["name"]}.json')
    else:
        logging.error('声学模型json文件重名，请重命名')

def add_lexicon_dict():
    name='ld1'
    g2p_dict_file=''
    g2p_file='lexicon_dicts.ld1'
    if os.path.exists(g2p_file):
        if not os.path.abspath(g2p_file).startswith(os.path.abspath('lexicon_dicts')):
            target_g2p_py=f'lexicon_dicts/{os.path.splitext(g2p_file.split("/")[-1])[0]}.py'
            if os.path.basename(target_g2p_py) in os.listdir('lexicon_dicts'):
                logging.error('字素-音素转换模块重名')
            shutil.copy(g2p_file,target_g2p_py)
            g2p_module_name=target_g2p_py[:-3].replace('/','.')
        else:
            g2p_module_name=os.path.relpath(g2p_file,'.').replace(path_sep,'.')
            g2p_module_name=g2p_module_name if not g2p_module_name.endswith('.py') else g2p_module_name[:-3]
    elif g2p_file!='':
        g2p_module_name=g2p_file if not g2p_file.endswith('.py') else g2p_file[:-3]
        g2p_module_name=g2p_module_name.replace('/','.')
    else:
        logging.error('字素-音素转换模块名不能为空')
    g2p_classname='chinese_sentence2pinyin_sequence'
    g2p_parameters=''
    g2p_attribute=''
    if g2p_dict_file.endswith('.py'):
        if os.path.exists(g2p_dict_file):
            if not os.path.abspath(g2p_dict_file).startswith(os.path.abspath('lexicon_dicts')):
                target_g2p_dict_py=f'lexicon_dicts/{os.path.splitext(g2p_dict_file.split("/")[-1])[0]}.py'
                if os.path.basename(target_g2p_dict_py) in os.listdir('lexicon_dicts'):
                    logging.error('获取字素-音素对照表模块重名')
                shutil.copy(g2p_dict_file,target_g2p_dict_py)
                g2p_dict_module_name=target_g2p_dict_py[:-3].replace('/','.')
            else:
                g2p_dict_module_name=os.path.relpath(g2p_dict_file,'.')[:-3].replace(path_sep,'.')
        else:
            logging.error(f'{g2p_dict_file}不存在')
    elif g2p_dict_file!='' or not g2p_dict_file.endswith(('.txt','.json')):
        g2p_dict_module_name=g2p_dict_file.replace('/','.')
    else:
        g2p_dict_module_name=''

    g2p_dict_classname=''
    g2p_dict_parameters=''
    g2p_dict_attribute=''
    g2p={
        'dict_file':g2p_dict_file,
        'dict_modulename':g2p_dict_module_name,
        'dict_classname':g2p_dict_classname,
        'dict_parameters':g2p_dict_parameters,
        'dict_attribute':g2p_dict_attribute,
        'modulename':g2p_module_name,
        'classname':g2p_classname,
        'parameters':g2p_parameters,
        'attribute':g2p_attribute
    }
    penc_dict_file='lexicon_dicts/ld1.py'

    penc_file='lexicon_dicts/ld1.py'
    if os.path.exists(penc_file):
        if not os.path.abspath(penc_file).startswith(os.path.abspath('lexicon_dicts')):
            target_penc_py=f'lexicon_dicts/{os.path.splitext(penc_file.split("/")[-1])[0]}.py'
            if os.path.basename(target_penc_py) in os.listdir('lexicon_dicts'):
                logging.error('音素编码模块重名')
            shutil.copy(penc_file,target_penc_py)
            penc_module_name=target_penc_py[:-3].replace('/','.')
        else:
            penc_module_name=os.path.relpath(penc_file,'.').replace(path_sep,'.')
            penc_module_name=penc_module_name if not penc_module_name.endswith('.py') else penc_module_name[:-3]
    elif penc_file!='':
        penc_module_name=penc_file if not penc_file.endswith('.py') else penc_file[:-3]
        penc_module_name=penc_module_name.replace('/','.')
    else:
        logging.error('音素编码模块名不能为空')
    penc_classname='enc_sequence'
    penc_parameters=''
    penc_attribute=''

    if penc_dict_file.endswith('.py'):
        if os.path.exists(penc_dict_file):
            if not os.path.abspath(penc_dict_file).startswith(os.path.abspath('lexicon_dicts')):
                target_penc_dict_py=f'lexicon_dicts/{os.path.splitext(penc_dict_file.split("/")[-1])[0]}.py'
                if os.path.basename(target_penc_dict_py) in os.listdir('lexicon_dicts'):
                    logging.error('音素编码模块重名')
                shutil.copy(penc_dict_file,target_penc_dict_py)
                penc_dict_module_name=target_penc_dict_py[:-3].replace('/','.')
            else:
                penc_dict_module_name=os.path.relpath(penc_dict_file,'.')[:-3].replace(path_sep,'.')
        else:
            logging.error(f'{penc_dict_file}不存在')
    elif penc_dict_file!='' or not penc_dict_file.endswith(('.txt','.json')):
        penc_dict_module_name=penc_dict_file.replace('/','.')
    else:
        penc_dict_module_name=''
    penc_dict_classname='load_pinyin_dict'
    penc_dict_parameters='"lexicon_dicts/chinese_pinyin_dict.txt"'
    penc_dict_attribute=''
    penc={
        'dict_file':penc_dict_file,
        'dict_modulename':penc_dict_module_name,
        'dict_classname':penc_dict_classname,
        'dict_parameters':penc_dict_parameters,
        'dict_attribute':penc_dict_attribute,
        'modulename':penc_module_name,
        'classname':penc_classname,
        'parameters':penc_parameters,
        'attribute':penc_attribute
    }

    pdec_dict_file='lexicon_dicts/ld1.py'
    pdec_file='lexicon_dicts.ld1'
    if os.path.exists(pdec_file):
        if not os.path.abspath(pdec_file).startswith(os.path.abspath('lexicon_dicts')):
            target_pdec_py=f'lexicon_dicts/{os.path.splitext(pdec_file.split("/")[-1])[0]}.py'
            if os.path.basename(target_pdec_py) in os.listdir('lexicon_dicts'):
                logging.error('音素解码模块重名')
            shutil.copy(pdec_file,target_pdec_py)
            pdec_module_name=target_pdec_py[:-3].replace('/','.')
        else:
            pdec_module_name=os.path.relpath(pdec_file,'.').replace(path_sep,'.')
            pdec_module_name=pdec_module_name if not pdec_module_name.endswith('.py') else pdec_module_name[:-3]
    elif pdec_file!='':
        pdec_module_name=pdec_file if not pdec_file.endswith('.py') else pdec_file[:-3]
    else:
        logging.error('音素解码模块名不能为空')

    pdec_classname='dec_sequence'
    pdec_parameters=''
    pdec_attribute=''

    if pdec_dict_file.endswith('.py'):
        if os.path.exists(pdec_dict_file):
            if not os.path.abspath(pdec_dict_file).startswith(os.path.abspath('lexicon_dicts')):
                target_pdec_dict_py=f'lexicon_dicts/{os.path.splitext(pdec_dict_file.split("/")[-1])[0]}.py' # 注意file的路径必须以“/”分隔，如果路径分隔符与平台相关，需要按path_sep分隔
                if os.path.basename(target_pdec_dict_py) in os.listdir('lexicon_dicts'):
                    logging.error('获取音素解码字典模块重名')
                shutil.copy(pdec_dict_file,target_pdec_dict_py)
                pdec_dict_module_name=target_pdec_dict_py[:-3].replace('/','.')
            else:
                pdec_dict_module_name=os.path.relpath(pdec_dict_file,'.')[:-3].replace(path_sep,'.')
        else:
            logging.error(f'{pdec_dict_file}不存在')
    elif pdec_dict_file!='' or not pdec_dict_file.endswith(('.txt','.json')):
        pdec_dict_module_name=pdec_dict_file.replace('/','.')
    else:
        pdec_dict_module_name=''

    pdec_dict_classname='load_id2pinyin_dict'
    pdec_dict_parameters='"lexicon_dicts/chinese_pinyin_dict.txt"'
    pdec_dict_attribute=''
    pdec={
        'dict_file':pdec_dict_file,
        'dict_modulename':pdec_dict_module_name,
        'dict_classname':pdec_dict_classname,
        'dict_parameters':pdec_dict_parameters,
        'dict_attribute':pdec_dict_attribute,
        'modulename':pdec_module_name,
        'classname':pdec_classname,
        'parameters':pdec_parameters,
        'attribute':pdec_attribute
    }

    p2g_dict_file='lexicon_dicts.ld1'
    p2g_file='lexicon_dicts.ld1'
    if os.path.exists(p2g_file):
        if not os.path.abspath(p2g_file).startswith(os.path.abspath('lexicon_dicts')):
            target_p2g_py=f'lexicon_dicts/{os.path.splitext(p2g_file.split("/")[-1])[0]}.py'
            if os.path.basename(target_p2g_py) in os.listdir('lexicon_dicts'):
                logging.error('音素-字素转换模块重名')
            shutil.copy(p2g_file,target_p2g_py)
            p2g_module_name=target_g2p_py[:-3].replace('/','.')
        else:
            p2g_module_name=os.path.relpath(p2g_file,'.').replace(path_sep,'.')
            p2g_module_name=p2g_module_name if not p2g_module_name.endswith('.py') else p2g_module_name[:-3]
    elif p2g_file!='':
        p2g_module_name=p2g_file if not p2g_file.endswith('.py') else p2g_file[:-3]
    else:
        logging.error('音素-字素转换模块名不能为空')
    p2g_classname='pinyin2chinese_word'
    p2g_parameters=''
    p2g_attribute=''
    if p2g_dict_file.endswith('.py'):
        if os.path.exists(p2g_dict_file):
            if not os.path.abspath(p2g_dict_file).startswith(os.path.abspath('lexicon_dicts')):
                target_p2g_dict_py=f'lexicon_dicts/{os.path.splitext(p2g_dict_file.split("/")[-1])[0]}.py'
                if os.path.basename(target_p2g_dict_py) in os.listdir('lexicon_dicts'):
                    logging.error('获取音素-字素对照表模块重名')
                shutil.copy(p2g_dict_file,target_p2g_dict_py)
                p2g_dict_module_name=target_p2g_dict_py[:-3].replace('/','.')
            else:
                p2g_dict_module_name=os.path.relpath(p2g_dict_file,'.')[:-3].replace(path_sep,'.')
        else:
            logging.error(f'{p2g_dict_file}不存在')
    elif p2g_dict_file!='' or not p2g_dict_file.endswith(('.txt','.json')):
        p2g_dict_module_name=p2g_dict_file.replace('/','.')
    else:
        p2g_dict_module_name=''
    p2g_dict_classname='load_pinyin2chinese_word'
    p2g_dict_parameters='"lexicon_dicts/dict.txt"'
    p2g_dict_attribute=''

    p2g={
        'dict_file':p2g_dict_file,
        'dict_modulename':p2g_dict_module_name,
        'dict_classname':p2g_dict_classname,
        'dict_parameters':p2g_dict_parameters,
        'dict_attribute':p2g_dict_attribute,
        'modulename':p2g_module_name,
        'classname':p2g_classname,
        'parameters':p2g_parameters,
        'attribute':p2g_attribute
    }

    lexicon_dict_json={
        'name':name,
        'g2p':g2p,
        'penc':penc,
        'pdec':pdec,
        'p2g':p2g
    }

    if not f'{lexicon_dict_json["name"]}.json' in os.listdir('jsons/lexicon_dicts'):
        dump_json(lexicon_dict_json,f'jsons/lexicon_dicts/{lexicon_dict_json["name"]}.json')
    else:
        logging.error('发音字典json文件重名，请重命名')

def add_language_model():
    name='2gram'
    typ='ngram'
    filetype='source'
    content='structure+weights'
    weights_file=''
    weights_type=''
    file='lm/gram2.py'
    if filetype=='source':
        if os.path.exists(file):
            if not os.path.abspath(file).startswith(os.path.abspath('lm')):
                target_lm_py=f'lm/{os.path.splitext(file.split("/")[-1])[0]}.py'
                if os.path.basename(target_lm_py) in os.listdir('lm'):
                    logging.error('语言模型模块重名')
                shutil.copy(file,target_lm_py)
                modulename=target_lm_py[:-3].replace('/','.')
            else:
                modulename=os.path.relpath(file,'.').replace(path_sep,'.')
                modulename=modulename if modulename.endswith('.py') else modulename[:-3]
        elif file=='':
            modulename=f'lm.{name}'
        else:
            modulename=file if not file.endswith('.py') else file[:-3]
            modulename=modulename.replace('/','.')
    classname='return_gram2_chinese_func'
    parameters=''
    attribute=''

    lm_json={
        'name':name,
        'type':typ,
        'filetype':filetype,
        'content':content,
        'weights_file':weights_file,
        'weights_type':weights_type,
        'file':file
    }

    if filetype=='source':
        lm_json['modulename']=modulename
        lm_json['classname']=classname
        lm_json['parameters']=parameters
        lm_json['attribute']=attribute
    
    if not f'{lm_json["name"]}.json' in os.listdir('jsons/lms'):
        dump_json(lm_json,f'jsons/lms/{lm_json["name"]}.json')
    else:
        logging.error('语言模型json文件重名，请重命名')

def add_decoder():
    name='chinese_pinyin2word_2gram_dec'
    file='decoders/chinese_pinyin2word_decoder.py'
    if os.path.exists(file):
        if not os.path.abspath(file).startswith(os.path.abspath('decoders')):
            target_decoder_py=f'decoders/{os.path.splitext(file.split("/")[-1])[0]}.py'
            if os.path.basename(target_decoder_py) in os.listdir('decoders'):
                logging.error('解码器模块重名')
            shutil.copy(file,target_decoder_py)
            modulename=target_decoder_py[:-3].replace('/','.')
        else:
            modulename=os.path.relpath(file,'.').replace(path_sep,'.')
            modulename=modulename if not modulename.endswith('.py') else modulename[:-3]
    elif file=='':
        modulename=f'decoders.{name}'
    else:
        modulename=file if not file.endswith('.py') else file[:-3]
        modulename=modulename.replace('/','.')
    classname='decode_chinese_pinyin2word_2gram'
    parameters=''
    attribute=''
    decoder_json={
        'name':name,
        'modulename':modulename,
        'classname':classname,
        'parameters':parameters,
        'attribute':attribute
    }
    if not f'{decoder_json["name"]}.json' in os.listdir('jsons/decoders'):
        dump_json(decoder_json,f'jsons/decoders/{decoder_json["name"]}.json')
    else:
        logging.error('解码器json文件重名，请重名')

def add_model():
    add_acoustic_model()
    add_lexicon_dict()
    add_language_model()
    add_decoder()

if __name__=='__main__':
    load_logging_config()
    add_lexicon_dict()